from v3dsolver.tree_manager.vel_calculator import VelCalculator
from v3dsolver.tree_manager.distribution_manager import get_domain
from v3dsolver.tree_manager.tree_manager import TreeManager
from numpy import array
from mpi4py import MPI
import time

class AdvectionManager:
    """
    Generic class for performing advection of elements.
    """
    def __init__(self,elems,dt,freestream,npmax=None,npseudo=None,
                 blob_rad=None,vel_method='d'):
        """
        Initialization of AdvectionManager class. Given a list of elems and 
        timestep dt, this class is responsible for the advection of all of 
        them.
        
        npmax: splitting criteria for tree
        npseudo: number of pseudo particles to generate per cell
        blob_rad: radius of blobs
        vel_method: velocity computation method to use:
                   'd': direct velocity computation
                   's': serial tree velocity computation
                   'p': parallel tree velocity computation (default)
        """
        self.elems = elems
        self.dt = dt
        self.freestream = freestream
        self.vel_method = vel_method
        self.npmax = npmax
        self.npseudo = npseudo
        self.blob_rad = blob_rad
    
    def compute_vels(self):
        """
        Computes velocity on all the elems, using direct computation, serial 
        tree computation or parallel computation depending on the parameter
        self.vel_method.
        """
        if self.vel_method=='d':
            VelCalculator().direct_vel_computation(self.elems)
            for elem in self.elems:
                if elem.is_effect():
                    elem.vel += self.freestream.get_vel(elem)
        elif self.vel_method=='s':
            domain = get_domain(self.elems)
            tree = TreeManager(self.elems,domain,self.npmax,self.npseudo,
                               self.blob_rad,comm=MPI.COMM_WORLD)
            VelCalculator(tree).serial_vel_computation()
            for elem in self.elems:
                if elem.is_effect():
                    elem.vel += self.freestream.get_vel(elem)
        elif self.vel_method=='p':
            domain = get_domain(self.elems,findglobal=True)
            tree = TreeManager(self.elems,domain,self.npmax,self.npseudo,
                               self.blob_rad,comm=MPI.COMM_WORLD)
            time1 = time.time()
            VelCalculator(tree).parallel_vel_computation()
            time2 = time.time()
            print 'Time taken for', len(self.elems), 'particles:', time2-time1
            for elem in self.elems:
                if elem.is_effect():
                    elem.vel += self.freestream.get_vel(elem)
    
    def integrate_Euler(self,dt):
        """
        Performs Euler integration with timestep dt and pre-known velocity
        """
        for elem in self.elems:
            if elem.is_effect():
                elem.pos += elem.vel*dt
    
    def move_Euler(self):
        """
        Computes velocity on all effect elems and moves them accordingly using
        Euler integration.
        """
        self.compute_vels()
        self.integrate_Euler(self.dt)
    
    def move_RK2(self):
        """
        Computes velocity on all effect elems and moves them accordingly using
        Runge Kutta order 2 integration scheme.
        """
        self.compute_vels()
        # save initial positions and velocities
        y0,k1,effect_elems = [],[],[]
        for i in range(len(self.elems)):
            elem = self.elems[i]
            if elem.is_effect():
                effect_elems += [elem]
                y0 += [elem.pos]
                k1 += [elem.vel]
        # move first step of two in RK2 method
        self.integrate_Euler(self.dt)
        # compute new velocities for next step
        self.compute_vels()
        # save new velocities
        k2 = []
        for elem in effect_elems:
            k2 += [elem.vel]
        # update positions using average of k1 and k2
        for i in range(len(effect_elems)):
            elem = effect_elems[i]
            elem.pos = y0[i] + 0.5*(k1[i]+k2[i])*self.dt
